;; CHANNELS

;; channel-receive: The channel, which is used by the work-distributor to
;; receive messages from the outer context. This channel is created in the
;; pool-initializer and given to all workers, so that they can report themselves
;; as ready to receive more work.

;; channel-return: This channel is created by the publish procedure and given to
;; work-distributor, which gives it to the workers, so that they can send
;; results of completed work as messages on this channel to the outer context.

(define-module (fibers-pool))


(use-modules
 ;; FIFO queue, not functional, using mutation
 ;; https://www.gnu.org/software/guile/manual/html_node/Queues.html
 (ice-9 q)
 (ice-9 match)
 (ice-9 threads)
 (rnrs exceptions)
 (rnrs conditions)
 ;; fibers internals are needed for creating schedulers without running anything
 ;; in them immediately
 (fibers)
 (fibers channels)
 (fibers internal))


(define displayln
  (lambda (msg)
    (display msg)
    (newline)))


(define run-worker-with-work
  (lambda (scheduler work-thunk-and-channel-return)
    ;; TODO: Somehow asynchronously start fibers as workers, avoiding the
    ;; overhead of one thread per worker, if possible.
    (call-with-new-thread
     (lambda ()
       (run-fibers
        (lambda ()
          (spawn-fiber (lambda () (worker index channel-receive))))
        #:scheduler scheduler)))))

(define work-distributor
  (lambda* (scheduler channel-receive #:key (max-busy-workers-count +inf.0))
    (let loop ([busy-workers-count 0]
               [work-queue (make-q)])
      (displayln "[WORK-DISTRIBUTOR]: work-distributor is listening for messages")

      (display "[WORK-DISTRIBUTOR]: number of works in queue: ")
      (displayln (q-length work-queue))

      ;; Get a message form the channel-receive and match it against expected
      ;; cases: (1) new work to be done (2) worker finished.
      (match (pk 'work-distributor-received-msg (get-message channel-receive))
        ;; A worker has finished its work. If there is more work in the work
        ;; queue, start a new worker to work on that work, otherwise loop again
        ;; with reduced busy workers count for the finished worker, to read the
        ;; next message, if there is any.
        [('worker-finished)
         (cond
          [(q-empty? work-queue)
           (loop (- busy-workers-count 1) work-queue)]
          [else
           ;; Invariant: When a worker finishes, the maximum of simultaneously
           ;; busy workers cannot be reached, because the worker was one of
           ;; the workers, which were started when the number of
           ;; simultaneously busy workers was lower than the maximum,
           ;; otherwise the worker should never have been started.
           (run-worker-with-work scheduler (deq! work-queue))
           (loop (+ busy-workers-count 1) work-queue)])]

        ;; New work has been received. We check, whether the busy worker count
        ;; allows us to run more workers. If the maximum of simultaneously busy
        ;; workers is not yet reached, we start a new worker, which will be
        ;; given the received work. If the maximum of simultaneously busy
        ;; workers is reached, we put the received work in a queue to be worked
        ;; on later. Received work consists of a thunk to be run in a worker and
        ;; a channel, on which the worker shall report the result.
        [('work . work-thunk-and-channel-return)
         (cond
          [(< busy-workers-count max-busy-workers-count)
           (run-worker-with-work scheduler work-thunk-and-channel-return)
           (loop (+ busy-workers-count 1) work-queue)]
          [else
           (enq! work-queue work-thunk-and-channel-return)
           (loop busy-workers-count work-queue)])]

        ;; On any other message raise a condition. Do not loop, instead return.
        [other
         (raise
          (condition
           (make-error)
           (make-message-condition "work-distributor received unrecognized message")
           (make-irritants-condition (list other))))]))))


(define worker
  (lambda (worker-index channel-receive)
    (let ([channel-worker (make-channel)])
      (let loop ()
        ;; Report as ready. Give my own channel to the work-distributor to let
        ;; it send me work.
        (put-message channel-receive
                     (cons 'worker-ready
                           channel-worker))
        ;; Get messages sent to me by the work distributor on my own channel.
        (match (pk 'worker-got-msg (get-message channel-worker))
          ;; If I receive work, do the work and return it on the
          ;; channel-return. Afterwards report to the work-distributor that this
          ;; worker has finised its work on channel-receive. Do not loop, as
          ;; there will be a new worker created for each new work.
          [('work . (thunk . channel-return))
           (put-message channel-return (thunk))
           (put-message channel-receive 'worker-finished)]

          ;; On any other message raise a condition.
          [other
           (raise
            (condition
             (make-error)
             (make-message-condition "worker received unrecognized message")
             (make-irritants-condition (list other))))])))))


(define pool-initializer
  (lambda* (#:key (parallelism (current-processor-count)))
    (let ([channel-receive (make-channel)]
          [scheduler (make-scheduler #:parallelism parallelism)])
      (displayln "[POOL INIT]: will start work-distributor")
      (call-with-new-thread
       (lambda ()
         (work-distributor scheduler channel-receive)))
      ;; Return the channel for receiving work, so that the outside context can
      ;; make use of it when calling ~publish~ to publish work.
      channel-receive)))


(define publish
  (lambda (work-as-thunk channel-receive)
    ;; The result of the computation can be taken from ~channel-return~.
    (let ([channel-return (make-channel)])
      ;; Put work tagged as work on the receive channel of the work-distributor.
      (let ([work-message (cons 'work (cons work-as-thunk channel-return))])
        (display
         (simple-format
          #f "[PUBLISHER]: will publish the following work: ~a\n"
          work-message))
        (put-message channel-receive work-message))

      (displayln "[PUBLISHER]: work published")
      ;; Return the ~channel-return~, so that the outside context can get
      ;; results from it.
      channel-return)))


(define busy-work
  (lambda ()
    (let loop ([i 0])
      (cond
       [(< i 5e8) (loop (+ i 1))]
       [else i]))))


(define c-rec (pool-initializer #:parallelism 2))
(define c-ret-2 (publish (lambda () (busy-work)) c-rec))
(define c-ret-1 (publish (lambda () (busy-work)) c-rec))
(get-message c-ret-2)
(get-message c-ret-1)
